{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "042412d9-4f18-495a-affd-e0ba104ce703",
   "metadata": {},
   "source": [
    "# BM25代码实践\n",
    "Milvus版本>=2.4.0\n",
    "marshmallow<4.0\n",
    "\n",
    "如果不对，重新安装，并且一定要记得重启内核，重启内核，不然还会用旧的！"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b202a791-d203-4caa-99fb-abb74e413a75",
   "metadata": {
    "ExecutionIndicator": {
     "show": false
    },
    "execution": {
     "iopub.execute_input": "2025-09-08T02:50:06.662997Z",
     "iopub.status.busy": "2025-09-08T02:50:06.662802Z",
     "iopub.status.idle": "2025-09-08T02:50:11.360649Z",
     "shell.execute_reply": "2025-09-08T02:50:11.360142Z",
     "shell.execute_reply.started": "2025-09-08T02:50:06.662981Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mDEPRECATION: pytorch-lightning 1.7.7 has a non-standard dependency specifier torch>=1.9.*. pip 24.0 will enforce this behaviour change. A possible replacement is to upgrade to a newer version of pytorch-lightning or contact the author to suggest that they release a version with a conforming dependency specifiers. Discussion can be found at https://github.com/pypa/pip/issues/12063\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "spacy 3.7.0 requires typer<0.10.0,>=0.3.0, but you have typer 0.16.0 which is incompatible.\n",
      "pai-easycv 0.11.7 requires timm==0.5.4, but you have timm 1.0.19 which is incompatible.\n",
      "ms-swift 3.7.1 requires peft<0.17,>=0.11, but you have peft 0.17.0 which is incompatible.\n",
      "weasel 0.3.4 requires typer<0.10.0,>=0.3.0, but you have typer 0.16.0 which is incompatible.\u001b[0m\u001b[31m\n",
      "\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.3.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m25.2\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install \"marshmallow<4.0\" --force-reinstall --quiet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "967858e6-b78a-4b81-9ad7-dcc755a65826",
   "metadata": {
    "ExecutionIndicator": {
     "show": false
    },
    "execution": {
     "iopub.execute_input": "2025-09-08T02:50:13.017226Z",
     "iopub.status.busy": "2025-09-08T02:50:13.016990Z",
     "iopub.status.idle": "2025-09-08T02:50:21.916584Z",
     "shell.execute_reply": "2025-09-08T02:50:21.916085Z",
     "shell.execute_reply.started": "2025-09-08T02:50:13.017209Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mDEPRECATION: pytorch-lightning 1.7.7 has a non-standard dependency specifier torch>=1.9.*. pip 24.0 will enforce this behaviour change. A possible replacement is to upgrade to a newer version of pytorch-lightning or contact the author to suggest that they release a version with a conforming dependency specifiers. Discussion can be found at https://github.com/pypa/pip/issues/12063\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "spacy 3.7.0 requires typer<0.10.0,>=0.3.0, but you have typer 0.16.0 which is incompatible.\n",
      "ms-swift 3.7.1 requires peft<0.17,>=0.11, but you have peft 0.17.0 which is incompatible.\u001b[0m\u001b[31m\n",
      "\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.3.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m25.2\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install pymilvus==2.4.7 torch --quiet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6f68ad87-f798-4172-96d2-4eee01d9ded6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-09-08T02:50:24.783011Z",
     "iopub.status.busy": "2025-09-08T02:50:24.782804Z",
     "iopub.status.idle": "2025-09-08T02:50:30.284261Z",
     "shell.execute_reply": "2025-09-08T02:50:30.283754Z",
     "shell.execute_reply.started": "2025-09-08T02:50:24.782994Z"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33mDEPRECATION: pytorch-lightning 1.7.7 has a non-standard dependency specifier torch>=1.9.*. pip 24.0 will enforce this behaviour change. A possible replacement is to upgrade to a newer version of pytorch-lightning or contact the author to suggest that they release a version with a conforming dependency specifiers. Discussion can be found at https://github.com/pypa/pip/issues/12063\u001b[0m\u001b[33m\n",
      "\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.3.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m25.2\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install \"pymilvus[model]\" --quiet"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18234633-fdf8-41bc-9837-118fad52a547",
   "metadata": {},
   "source": [
    "假设下面每一个字符串都是一个文档"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1d7305cd-3775-4155-aeb0-bb5e9990ee6c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-09-03T08:42:42.466055Z",
     "iopub.status.busy": "2025-09-03T08:42:42.465841Z",
     "iopub.status.idle": "2025-09-03T08:42:42.470698Z",
     "shell.execute_reply": "2025-09-03T08:42:42.470307Z",
     "shell.execute_reply.started": "2025-09-03T08:42:42.466040Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "docs = [\n",
    "    \"机器学习正在改变我们的生活方式。\",\n",
    "    \"深度学习在图像识别中表现出色。\",\n",
    "    \"自然语言处理是计算机科学的重要领域。\",\n",
    "    \"自动驾驶依赖于先进的算法。\",\n",
    "    \"AI可以帮助医生诊断疾病。\",\n",
    "    \"金融领域广泛应用数据分析技术。\",\n",
    "    \"生产效率可以通过自动化技术提高。\",\n",
    "    \"机器智能的未来充满潜力。\",\n",
    "    \"大数据支持是机器智能发展的关键。\",\n",
    "    \"量子隧穿效应使得电子能够穿过经典力学认为无法穿过的势垒，这在半导体器件中有着重要的应用。\"\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b814302a-1d55-48f3-b076-b8388fce6f1a",
   "metadata": {},
   "source": [
    "使用BM25对第一个文档“机器学习正在改变我们的生活方式”做分词处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b85df11-b2d5-488e-8bd1-883f4e872cb4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from pymilvus.model.sparse.bm25.tokenizers import build_default_analyzer\n",
    "from pymilvus.model.sparse import BM25EmbeddingFunction\n",
    "\n",
    "# 使用支持中文的分析器\n",
    "analyzer = build_default_analyzer(language=\"zh\")\n",
    "\n",
    "# 分析器对文本做分词处理\n",
    "tokens1 = analyzer(docs[0])\n",
    "print(tokens1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c85572fb-dd18-4f31-922a-6dcb4dc87250",
   "metadata": {},
   "source": [
    "分词结果如下：\n",
    "\n",
    "['机器', '学习', '改变', '生活', '方式']\n",
    "\n",
    "接下来对整个文档做分词处理，并且计算文档集合的IDF等参数："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "179abe91-133b-4b6a-bd32-2b988eebcce9",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-09-03T08:44:14.051047Z",
     "iopub.status.busy": "2025-09-03T08:44:14.050654Z",
     "iopub.status.idle": "2025-09-03T08:44:14.070774Z",
     "shell.execute_reply": "2025-09-03T08:44:14.070345Z",
     "shell.execute_reply.started": "2025-09-03T08:44:14.051027Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 创建BM25EmbeddingFunction实例，传入分词器，以及其他参数\n",
    "bm25_ef = BM25EmbeddingFunction(analyzer)\n",
    "\n",
    "# 计算文档集合的参数\n",
    "bm25_ef.fit(docs)\n",
    "\n",
    "# 保存训练好的参数到磁盘以加快后续处理\n",
    "bm25_ef.save(\"bm25_params.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca3eceb2-9c20-46d0-b243-579a11ca81a3",
   "metadata": {},
   "source": [
    "查看参数具体内容"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "883815ac-df0b-4849-a9e0-378fd663bcfc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "file_path = \"bm25_params.json\"\n",
    "with open(file_path, 'r', encoding='utf-8') as file:\n",
    "    bm25_params = json.load(file)\n",
    "    print(bm25_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea98d7f7-d7bc-48ff-9e6b-157ad8f68257",
   "metadata": {},
   "source": [
    "corpus_size是文档数量，可以看到avgdl和idf_value两个参数"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8afcba14-33cf-44c9-80a3-6f24b31f4a0e",
   "metadata": {},
   "source": [
    "idf_word 是 BM25对文档集合的分词结果，也就是前面提到的词汇表。词汇表中单词的数量，也是稀疏向量的维度。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86ba2241-80d6-407c-bbc8-235843558f7e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# BM25词汇表中的单词数量\n",
    "print(f\"BM25词汇表中的单词数量：{len(bm25_params['idf_word'])}\")\n",
    "\n",
    "# BM25稀疏向量的维度\n",
    "print(f\"BM25稀疏向量维度：{bm25_ef.dim}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03a61757-c542-4cb3-84f9-2454514da183",
   "metadata": {},
   "source": [
    "返回的结果：\n",
    "\n",
    "BM25词汇表中的单词数量：46\n",
    "BM25稀疏向量维度：46\n",
    "\n",
    "需要的参数计算好了，接下来就可以生成文档集合的稀疏向量了。文档集合中有10篇文档，也就是10个字符串，而稀疏向量的维度是46，所以文档集合的稀疏向量是一个10行46列的矩阵。每一行表示一个文档的稀疏向量。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ad57c28-ba9a-4d02-a99e-915fea9efabc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 生成文档集合的稀疏向量\n",
    "sparse_vectors_bm25 = bm25_ef.encode_documents(docs)\n",
    "\n",
    "# 打印文档集合的稀疏向量\n",
    "print(sparse_vectors_bm25)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2371f913-0529-45b8-9feb-3b891839cff3",
   "metadata": {},
   "source": [
    "我们来看下第一个文档“机器学习正在改变我们的生活方式。”的稀疏向量："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2deeab0b-91be-468c-867b-9d0396456aa2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 第一个文档的稀疏向量\n",
    "print(list(sparse_vectors_bm25)[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f035918-37ad-4bde-9983-0d57dfdf5cc1",
   "metadata": {},
   "source": [
    "结果为：\n",
    "\n",
    "(0, 0)        1.0344827586206897\n",
    "  (0, 1)        1.0344827586206897\n",
    "  (0, 2)        1.0344827586206897\n",
    "  (0, 3)        1.0344827586206897\n",
    "  (0, 4)        1.0344827586206897\n",
    "你发现了吧，第一个文档的稀疏向量只有5个非零元素，因为它的分词结果是5个单词，对应上了。而且，每个元素的值都相同，说明它们的逆文档频率 IDF 和词频 TF 都是一样的。\n",
    "\n",
    "第一个文档的分词结果：\n",
    "\n",
    "['机器', '学习', '改变', '生活', '方式']\n",
    "文档集合处理好了，我们再给出一个查询的句子，就可以执行搜索了。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e2a82c9-cf8a-48ac-93ef-74a06cc419da",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "query = [\"自动驾驶如何影响汽车行业？\"]\n",
    "\n",
    "# 把查询文本向量化\n",
    "query_sparse_vectors_bm25 = bm25_ef.encode_queries(query)\n",
    "\n",
    "# 打印稀疏向量\n",
    "print(query_sparse_vectors_bm25)\n",
    "\n",
    "# 查询的分词结果\n",
    "print(analyzer(query[0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27d757c8-de7c-476f-95bb-96786b1c55a1",
   "metadata": {},
   "source": [
    "查看查询的稀疏向量，以及它的分词结果。\n",
    "\n",
    "(0, 12)        1.845826690498331\n",
    "  (0, 13)        1.845826690498331\n",
    "['自动', '驾驶', '影响', '汽车行业']\n",
    "你可能会有疑问，为什么查询分词后得到4个单词，但是它的稀疏向量只有2维？因为这4个单词中，词汇表中只有“自动”和“驾驶”，没有“影响”和“汽车行业”，后两个词的 BM25 分数为0。\n",
    "\n",
    "哎，毕竟是门外汉啊。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
